[PyTorch]Fused ScaledActivation Kernel in TE Ops and reducing CPU overheads for MOE - #3238
Conversation
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
for more information, see https://pre-commit.ci
Greptile SummaryThis PR introduces fused C++ kernels for scaled activations (SReLU, SwiGLU, clamped SwiGLU), replaces the multi-step Python-level arithmetic in
Confidence Score: 5/5Safe to merge — the changes are internally consistent refactors with no correctness regressions identified. All forward/backward offset index adjustments are consistent across forward save and backward unpack sites. The new fused C++ kernels preserve the mathematical semantics of the old multi-step Python paths (scale multiply, activation, interleave reshape). The shared-expert (num_groups==1) path correctly sets new offset tensors to None, and the branch guard ensures the nvfp4 dgrad path that consumes fc1_x_tensor_offsets is never reached for that case. The only dead-code guard ( Files Needing Attention: No files require special attention. Important Files Changed
Sequence DiagramsequenceDiagram
participant PY as Python (fuser_forward)
participant TEX as tex (C++ ext)
participant KERNEL as CUDA Kernel
Note over PY,KERNEL: Forward Pass
PY->>TEX: "splits_to_offsets_multi(split_sizes, strides=[...])"
TEX-->>PY: split_sizes, [base_offsets, in_offsets, out_offsets, ...]
PY->>TEX: "group_quantize(x, quantizer, num_groups, split_sizes, tensor_offsets=in_offsets)"
TEX->>KERNEL: fused quantize+group
KERNEL-->>PY: grouped_x
PY->>TEX: scaled_swiglu(input, scales, quantizer, glu_interleave_size)
TEX->>KERNEL: nvte_scaled_swiglu (fused act+scale)
KERNEL-->>PY: activation_out
Note over PY: save [split_sizes, base_offsets, split_points, in_offsets, out_offsets, ...]
Note over PY,KERNEL: Backward Pass (reuses saved offsets)
PY->>TEX: "group_quantize(dy, quantizer, N, split_sizes, tensor_offsets=out_offsets)"
TEX->>KERNEL: fused quantize+group
KERNEL-->>PY: grouped_dy
PY->>TEX: scaled_dswiglu(grad, input, scales, quantizer, glu_interleave_size, compute_scale_grad)
TEX->>KERNEL: nvte_scaled_dswiglu (fused grad_input + grad_scales)
KERNEL-->>PY: grad_input, grad_scales
Reviews (16): Last reviewed commit: "Merge branch 'main' into grouped_linear_..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
for more information, see https://pre-commit.ci
| ) | ||
|
|
||
|
|
||
| class ForwardScaledActivationGroupedLinear(FusedOperation): |
There was a problem hiding this comment.
There is a trade-off in creating a partially fused module, since we already have FC1 - ACT - FC2 fused module, We need to justify the value of creating a FC1-ACT fusion, instead of just adding features to the grouped_mlp instead.
There was a problem hiding this comment.
This essentially enables the infrastructure to enable (Act + GroupQuant from FC2) fusion, and enables to add fused kernels if possible in the future. cc: @timmoon10
There was a problem hiding this comment.
The other important thought i had in mind is the fact that ScaledActivation today doesnt take in m_splits and it is going to be a lot of upstream disruption to allow for that. And at the same time ScaledActivation can be used after Dense Layers as well and not necessarily after a GroupedLinear layer(and so it might not always need m_splits).
Allowing for this fusion, we allow the m_splits information to be also consumed in the scaled_activation + grouped quantization fusion. Right now the activation kernel that we are using isnt even using the m_splits, but for paged stashing optimization we might need that and we can potentially add a new kernel for that in the future.
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
for more information, see https://pre-commit.ci
|
/te-ci pytorch |
…ns for precomputed tensor offsets in backward Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
for more information, see https://pre-commit.ci
vthumbe1503
changed the title
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
vthumbe1503
changed the title
…ansformerEngine into grouped_linear_act_fusion
|
/te-ci pytorch |
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci ugly solution for dbias fusion Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci
872d964 to
802fb46
Compare
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
for more information, see https://pre-commit.ci
vthumbe1503
changed the title
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
vthumbe1503
changed the title
|
/te-ci pytorch |
|
Pipeline: 61451302 |
4 participants
Description
Better Kernels for Scaled Activations
Precomputed Tensor Offsets Optimization
Also strengthens test_grouped_mlp assertions. We expect fusion even for nvfp4_rht + scale_swiglu + glu_interleave_size = 32 case.
BF16 Grouped MLP Performance for num_groups=8 and swiglu activation
MXFP8 Grouped MLP performance for num_groups = 8 and swiglu activation
With CuteDSL fusions
Without CuteDSL fusions
Fixes #2988
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: